import os
import codecs
from metrics.iou import DetectionIoUMetrics

base_dir = "/data/ppocr/trainData/test"
gt_dir = os.path.join(base_dir, "gt")
pred_dir = os.path.join(base_dir, "ydbs")
gt_lst = []
pred_lst = []
for name in os.listdir(gt_dir):
    gt_line = []
    pred_line = []
    with codecs.open(os.path.join(gt_dir, name), "r", "utf8") as f:
        for gt in f.readlines():
            ids = [int(idx) for idx in gt.split(",")]
            gt_line.append([(ids[0], ids[1]), (ids[2], ids[3]), (ids[4], ids[5]), (ids[6], ids[7])])

    with codecs.open(os.path.join(pred_dir, name), "r", "utf8") as f:
        for pred in f.readlines():
            ids = [int(float(idx)) for idx in pred.strip().split(",")[0:8]]
            pred_line.append([(ids[0], ids[1]), (ids[2], ids[3]), (ids[4], ids[5]), (ids[6], ids[7])])
    pred_lst.append(pred_line)
    gt_lst.append(gt_line)

results = []
evaluator = DetectionIoUMetrics()
for gt1, pred1 in zip(gt_lst, pred_lst):
    results.append(evaluator.evaluate_image(gt1, pred1))
metrics = evaluator.combine_results(results)


# 120epoch {'precision': 0.8254207141879875, 'recall': 0.7132048705520747, 'hmean': 0.7652207001522069}
# baidu iou threshold=0.5  {'precision': 0.7169438962503524, 'recall': 0.6012531032036884, 'hmean': 0.654021732141709}
# zhangbo {'precision': 0.853786315288141, 'recall': 0.6357725499468022, 'hmean': 0.7288250440439085}
# {'precision': 0.8328133125325013, 'recall': 0.7572999172479017, 'hmean': 0.7932635750108353}
# ydbs {'precision': 0.8557384545038866, 'recall': 0.8849745832840761, 'hmean': 0.8701110013366653}

print(metrics)
